package com.hanxry.filter;

import cn.hutool.core.util.StrUtil;
import org.slf4j.MDC;
import org.springframework.stereotype.Component;

import javax.servlet.*;
import javax.servlet.annotation.WebFilter;
import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
import java.util.UUID;

@Component
@WebFilter(
        filterName = "traceFilter",
        urlPatterns = {"/*"},
        asyncSupported = true)
public class TraceFilter implements Filter {
    public TraceFilter() {
    }

    @Override
    public void init(FilterConfig filterConfig) {
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws ServletException, IOException {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        // 从请求头获取trace_id，如果有，则加入其中
        String xxTraceId = httpRequest.getHeader("xx-trace_id");
        if (StrUtil.isEmpty(xxTraceId)) {
            xxTraceId = UUID.randomUUID().toString().replaceAll("-", "");
        }
        MDC.put("traceId", xxTraceId);
        // 分布式调用时，可以记录rpcId
        String xxRpcId = httpRequest.getHeader("xxRpcId");
        if (StrUtil.isNotEmpty(xxRpcId)) {
            MDC.put("rpcId", xxRpcId);
        }
        // 记录客户端 IP
        MDC.put("clientIp", this.getClientIp(httpRequest));

        chain.doFilter(request, response);
    }

    @Override
    public void destroy() {
        MDC.remove("traceId");
        MDC.remove("rpcId");
        MDC.remove("clientIp");
    }

    private String getClientIp(HttpServletRequest request) {
        if (request == null) {
            return "";
        } else {
            String ip = request.getHeader("x-forwarded-for");
            if (StrUtil.isEmpty(ip)) {
                ip = request.getHeader("Proxy-Client-IP");
            }

            if (StrUtil.isEmpty(ip)) {
                ip = request.getHeader("WL-Proxy-Client-IP");
            }

            if (StrUtil.isEmpty(ip)) {
                ip = request.getHeader("X-Real-IP");
            }

            if (StrUtil.isEmpty(ip)) {
                ip = request.getRemoteAddr();
            }

            if (ip != null && ip.contains(",")) {
                ip = ip.split(",")[0];
            }

            if (ip != null) {
                int index = ip.indexOf(",");
                ip = index > -1 ? ip.substring(0, index) : ip;
            }
            return "0:0:0:0:0:0:0:1".equals(ip) ? "127.0.0.1" : ip;
        }
    }
}
